import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
# 修复matplotlib中文显示问题
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'sans-serif']
matplotlib.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'sans-serif']
plt.rcParams['axes.unicode_minus'] = False

import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.preprocessing import StandardScaler

# 导入所有模型
from trainCode.RAE_GCN_Transformer_BSM1 import RAE_GCN_Transformer, BSM1Dataset
from trainCode.Transformer_BSM1 import Transformer
from trainCode.GCN_BSM1 import GCN
from trainCode.GCN_Transformer_BSM1 import GCN_Transformer
from trainCode.RAE_Transformer_BSM1 import RAE_Transformer

def load_model_results(model_name):
    """加载模型的评估结果"""
    try:
        acc_val = np.load(f'model/acc_val_{model_name}_BSM1.npy')
        return acc_val
    except FileNotFoundError:
        print(f"警告: 未找到模型 {model_name} 的验证准确率文件")
        return None

def evaluate_model(model, model_name, device):
    """评估模型在测试集上的性能"""
    print(f"\n正在评估 {model_name} 模型...")
    
    # 检查模型参数
    if hasattr(model, 'embed_dim') and hasattr(model, 'num_heads'):
        embed_dim = getattr(model, 'embed_dim')
        num_heads = getattr(model, 'num_heads')
        if embed_dim % num_heads != 0:
            print(f"警告: 模型 {model_name} 的 embed_dim({embed_dim}) 不能被 num_heads({num_heads}) 整除")
            print("这可能导致AssertionError: embed_dim must be divisible by num_heads错误")
            return None, None
        else:
            print(f"模型参数检查：embed_dim({embed_dim}) 可以被 num_heads({num_heads}) 整除")
    
    # 加载测试数据
    test_data = BSM1Dataset('BSM1_WWTP_data/test_data.csv', seq_len=10)
    test_loader = DataLoader(test_data, batch_size=512, shuffle=False)
    
    # 加载模型权重
    try:
        model.load_state_dict(torch.load(f'model/model_{model_name}_BSM1.pth', map_location=device))
        model.eval()
    except FileNotFoundError:
        print(f"错误: 未找到模型 {model_name} 的权重文件")
        return None, None
    except RuntimeError as e:
        print(f"错误: 无法加载模型 {model_name} 的权重，可能是模型结构与保存的权重不匹配")
        print(f"详细错误: {str(e)}")
        print(f"跳过模型 {model_name} 的评估")
        return None, None
    
    # 在测试集上评估
    try:
        y_true = []
        y_pred = []
        
        with torch.no_grad():
            for x, y in test_loader:
                x = x.to(device)
                y_true.extend(y.numpy())
                outputs = model(x)
                _, predicted = torch.max(outputs.data, 1)
                y_pred.extend(predicted.cpu().numpy())
    except Exception as e:
        print(f"错误: 模型 {model_name} 在推理过程中出错")
        print(f"详细错误: {str(e)}")
        return None, None
    
    # 计算准确率和混淆矩阵
    accuracy = accuracy_score(y_true, y_pred)
    cm = confusion_matrix(y_true, y_pred)
    class_accuracies = cm.diagonal() / cm.sum(axis=1)
    
    # 打印总体准确率
    print(f"{model_name} 测试集准确率: {accuracy:.4f}")
    
    # 打印每个类别的准确率
    print(f"\n{model_name} 在各故障类别上的准确率:")
    print("-" * 50)
    print("类别\t准确率")
    for i, acc in enumerate(class_accuracies):
        print(f"{i}\t{acc:.4f}")
    print(f"平均\t{accuracy:.4f}")
    print("-" * 50)
    
    return accuracy, class_accuracies

def compare_models():
    """比较所有模型的性能"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 获取数据集信息
    train_data = BSM1Dataset('BSM1_WWTP_data/train_data.csv', seq_len=10)
    input_dim = train_data.data.shape[1]
    num_classes = len(np.unique(train_data.labels))
    seq_len = 10
    
    # 定义所有模型（RAE_GCN_Transformer放到最后）
    models = {
        'Transformer': Transformer(input_dim=input_dim, seq_len=seq_len, num_classes=num_classes).to(device),
        'GCN': GCN(input_dim=input_dim, seq_len=seq_len, num_classes=num_classes).to(device),
        'GCN_Transformer': GCN_Transformer(input_dim=input_dim, seq_len=seq_len, num_classes=num_classes).to(device),
        'RAE_Transformer': RAE_Transformer(input_dim=input_dim, seq_len=seq_len, num_classes=num_classes).to(device),
        'RAE_GCN_Transformer': RAE_GCN_Transformer(input_dim=input_dim, seq_len=seq_len, num_classes=num_classes).to(device)
    }
    
    # 评估所有模型
    results = {}
    class_results = {}
    
    # 直接加载SVM和LDA的已保存结果，不重新训练
    print("\n加载SVM模型的已保存结果...")
    try:
        svm_acc = np.load('model/acc_test_SVM_BSM1.npy')[0]
        svm_class_acc = np.load('model/class_acc_SVM_BSM1.npy')
        results['SVM'] = svm_acc
        class_results['SVM'] = svm_class_acc
        print(f"SVM 测试集准确率: {svm_acc:.4f}")
    except FileNotFoundError:
        print("警告: 未找到SVM模型的测试准确率文件，请先运行 trainCode/SVM_BSM1.py 进行训练")
    
    print("\n加载LDA模型的已保存结果...")
    try:
        lda_acc = np.load('model/acc_test_LDA_BSM1.npy')[0]
        lda_class_acc = np.load('model/class_acc_LDA_BSM1.npy')
        results['LDA'] = lda_acc
        class_results['LDA'] = lda_class_acc
        print(f"LDA 测试集准确率: {lda_acc:.4f}")
    except FileNotFoundError:
        print("警告: 未找到LDA模型的测试准确率文件，请先运行 trainCode/LDA_BSM1.py 进行训练")
    
    # 评估其他深度学习模型
    for name, model in models.items():
        acc, class_acc = evaluate_model(model, name, device)
        if acc is not None:
            results[name] = acc
            class_results[name] = class_acc
    
    # 创建结果目录
    os.makedirs('results', exist_ok=True)
    
    # 打印每个故障类别的准确率
    print("\n各模型在每个故障类别上的准确率:")
    print("-" * 80)
    # 获取类别数量
    num_classes = len(next(iter(class_results.values())))
    
    # 打印表头
    header = "模型名称"
    for i in range(num_classes):
        header += f"\t{i}"
    header += "\t平均"
    print(header)
    
    # 打印每个模型的各类别准确率
    for model_name, class_acc in class_results.items():
        row = f"{model_name}"
        for i in range(num_classes):
            row += f"\t{class_acc[i]:.4f}"
        row += f"\t{results[model_name]:.4f}"
        print(row)
    print("-" * 80)
    
    # 生成表格4-4：不同模型对各故障及正常数据的分类效果
    generate_classification_table(results, class_results, num_classes)
    
    print("\n模型比较完成，表格4-4已保存到 'results' 目录")

def generate_classification_table(results, class_results, num_classes):
    """生成表格4-4：不同模型对各故障及正常数据的分类效果"""
    print("\n正在生成表格4-4：不同模型对各故障及正常数据的分类效果...")
    
    # 创建DataFrame用于保存表格数据
    # 设置列名：类别0-10和平均
    columns = [f'{i}' for i in range(num_classes)] + ['平均']
    
    # 创建空的DataFrame
    table_df = pd.DataFrame(columns=columns)
    
    # 定义模型显示顺序（RAE_GCN_Transformer放在最后）
    model_order = ['SVM', 'LDA', 'Transformer', 'GCN', 'GCN_Transformer', 'RAE_Transformer', 'RAE_GCN_Transformer']
    
    # 按照指定顺序填充数据（以百分比形式）
    for model_name in model_order:
        if model_name in class_results:
            class_acc = class_results[model_name]
            row_data = {}
            # 添加各类别准确率（百分比形式）
            for i in range(num_classes):
                row_data[f'{i}'] = f'{class_acc[i]*100:.2f}%'
            # 添加平均准确率
            row_data['平均'] = f'{results[model_name]*100:.2f}%'
            # 将行数据添加到DataFrame
            table_df.loc[model_name] = row_data
    
    # 重命名索引为"模型"
    table_df.index.name = '模型'
    
    # 保存为CSV文件
    try:
        # 确保results目录存在
        os.makedirs('results', exist_ok=True)
        
        csv_path = 'results/table4-4_classification_results.csv'
        table_df.to_csv(csv_path, encoding='utf-8-sig')  # 使用utf-8-sig编码以支持中文
        print(f"表格已保存到: {csv_path}")
    except PermissionError:
        print(f"警告: 无法保存表格到文件，可能是文件被占用或没有足够权限")
    except Exception as e:
        print(f"警告: 保存表格时出错: {str(e)}")
    
    # 打印表格内容
    print("\n表格4-4 不同模型对各故障及正常数据的分类效果")
    print("Table 4-4 Classification results of different models for normal data and faults")
    print(table_df.to_string())
    
    return table_df

if __name__ == '__main__':
    # 比较所有模型，使用已保存的结果
    compare_models()